import os, sys, shelve, collections, xmlrpclib, time, datetime
import cPickle as pickle
from SimpleXMLRPCServer import SimpleXMLRPCServer
import zlib
from bourbon.util import *
from bourbon.runner import JobRunner, RPCJobRunner
import tempfile
import subprocess as sp
		

class FileJobRunner(JobRunner):
	"""Runs a job from a file, and stores the results to a file."""
	def __init__(self, file_input, file_output):
		self._input = file_input
		self._output = file_output
		
	def start(self):
		# TODO: check to make sure input and output files are valid
		
		# load Experiment object from _input
		experiment = pickle.load(open(self._input, 'r'))
		
		time1 = datetime.datetime.now()
		# run the experiment
		self.run_experiment(experiment, steps=1)
		td = datetime.datetime.now() - time1
		duration = (td.microseconds + float(td.seconds + td.days * 24 * 3600) * 10**6) / 10**6
		results = experiment.results()
		
		# add duration to results
		results['Duration'] = duration
		
		# save results to file
		pickle.dump( results, open(self._output, 'w'))
		
	
class PBSJobRunner(RPCJobRunner):
	"""Extends RPC JobRunner to submit Bourbon jobs on a PBS job server."""
	def __init__(self, address="http://localhost:8000", max_jobs=10):
		RPCJobRunner.__init__(self, address)
		self._folder = '/shared.scratch/tr2/bourbon'
		self._input = 'input'
		self._output = 'output'
		self._error = 'error'
		self._results = 'results'
		self._scripts = 'scripts'
		self._conn = 'tr2@sugar.rice.edu'
		self._user = 'tr2'
		self._max_jobs = max_jobs
		
	def submit_to_pbs(self, j_id, job):
		"""Submit the job to the PBS Server."""
		jfile = "job_%d.bourbon"%j_id
		finput = os.path.join(self._folder, self._input, jfile)
		fresult = os.path.join(self._folder, self._results, jfile)
		fscript = os.path.join(self._folder, self._scripts, jfile)
		command = "bourbon runner file --input=%s --output=%s"%(finput, fresult)
		
		# write PBS script
		pbs = ""
		output = os.path.join(self._folder, self._output, 'job-%d.out'%j_id)
		opts = ['-q commons', '-l nodes=1:ppn=1,pmem=8000m,walltime=24:00:00', '-W x=NACCESSPOLICY:SINGLEJOB', '-o %s'%output, '-V', '-m n', '-j oe']
		
		for o in opts:
			pbs += "#PBS %s\n"%o
		
		pbs += "\n"
		pbs += command + "\n"
		
		
		tmp = tempfile.gettempdir()
		
		# submit pbs script
		jpbs = jfile
		scriptfile = os.path.join(tmp, jpbs)
		s_out = open(scriptfile, 'w')
		s_out.write(pbs)
		s_out.close()
		
		# copy file to server
		os.system("scp %s %s:%s"%(scriptfile, self._conn, fscript))
		
		os.unlink(scriptfile)
		
		# write job input data
		joutfile = os.path.join(tmp, jfile)
		jout = open(joutfile, 'w')
		pickle.dump(job, jout)
		jout.close()
		
		# scp file to server
		os.system("scp %s %s:%s"%(joutfile, self._conn, finput))
		
		# delete temp file
		os.unlink(joutfile)
		
		# submit the pbs script!
		print 'Submitting job %d'%j_id
		os.system('ssh %s "qsub %s"'%(self._conn, fscript))
		
		
	def retrieve_results(self):
		fresult = os.path.join(self._folder, self._results)
		#args = [self._conn,'"ls %s"'%fresult]
		cmd = 'ssh %s "ls %s"'%(self._conn, fresult)
		p1 = sp.Popen(cmd, shell=True, stdout=sp.PIPE, stderr=sp.PIPE)
		p1.wait()
		#files = p1.stdout.read()
		
		tmp = tempfile.gettempdir()
		
		self.verify_server()
		
		for afile in p1.stdout.readlines():
			afile = afile.strip()
			if afile.endswith('.bourbon'):
				j_id = int(afile[4:-8])
				# copy file to a local place
				the_file = os.path.join(fresult, afile)
				tmp_file = os.path.join(tmp, afile)
				jfile = afile
				finput = os.path.join(self._folder, self._input, jfile)
				#fresult = os.path.join(self._folder, self._results, jfile)
				fscript = os.path.join(self._folder, self._scripts, jfile)
				
				os.system("scp %s:%s %s"%(self._conn,the_file,tmp_file))
				
				# read contents
				result_data = open(tmp_file, 'r').read()
				
				# upload to server
				self._server.submit_results(j_id, result_data)
				
				# remove from server
				os.system('ssh %s "rm %s"'%(self._conn, the_file))
				os.system('ssh %s "rm %s"'%(self._conn, finput))
				os.system('ssh %s "rm %s"'%(self._conn,fscript))
				
				os.unlink(tmp_file)
				
	def num_running(self):
		cmd = 'ssh %s "showq | grep %s"'%(self._conn, self._user)
		p1 = sp.Popen(cmd, shell=True, stdout=sp.PIPE, stderr=sp.PIPE)
		p1.wait()
		
		num = 0
		for status in p1.stdout.readlines():
			status = status.strip()
			if len(status) > 0:
				num += 1
			
		return num
	
	def start(self):
		"""The PBSJobRunner will keep a certain number of jobs running on the PBS Server.  
			To accomplish this, it will periodically check the status of submitted jobs and submit the results of finished jobs.
		"""
		
		self.get_server()
		
		self.verify_server()
			
		while self._run:
			self.verify_server()
			
			# check for results
			self.retrieve_results()
			
			num = self.num_running()
			
			if num < self._max_jobs:
				for i in range( self._max_jobs - num):
					self.get_job()
					
			# sleep for a VERY LONG TIME (thirty min.)
			time.sleep(60.0*20)
		
		
	def get_job(self):		
		in_data = self._server.get_next_work(True)

		if in_data == 0: 
			self._run = False
			#time.sleep(5.0)
			pass
		else:
			e_id, exp_data = in_data
			print 'Got Job [%d], size [%d]'%(e_id, len(exp_data.data) )
			experiment = exp_from_data( exp_data.data )
			self.submit_to_pbs(e_id, experiment)

			del experiment
			del exp_data

		